"""
交易策略回测 Web 应用
功能：
1. 配置回测参数并运行
2. 实时查看回测进度和 Web UI
3. 查看历史回测结果
4. 对比不同参数的回测效果
"""
import os
import sys
import json
import subprocess
import threading
import sqlite3
from datetime import datetime, date
from flask import Flask, render_template, request, jsonify, Response
import pandas as pd

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from utils.strategy_recorder import StrategyDBRecorder

app = Flask(__name__)
app.config['SECRET_KEY'] = 'trading-backtest-secret-key'

# 全局变量存储当前运行的进程
current_process = None
current_output = []
current_webui_port = None

# 获取项目根目录
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
db_path = os.path.join(project_root, '..', 'trading_backtest.db')
db_recorder = StrategyDBRecorder(db_path=db_path)


# ==================== 路由：主页 ====================
@app.route('/')
def index():
    """主页：显示策略列表和配置表单"""
    return render_template('index.html')


# ==================== 路由：获取策略列表 ====================
@app.route('/api/strategies')
def get_strategies():
    """获取可用的策略列表（从数据库和文件系统）"""
    # 从数据库获取已运行过的策略
    conn = sqlite3.connect(db_path)
    query = '''
        SELECT DISTINCT strategy_name, COUNT(*) as count
        FROM backtest_results 
        GROUP BY strategy_name
        ORDER BY count DESC, strategy_name
    '''
    df = pd.read_sql_query(query, conn)
    conn.close()
    
    db_strategies = {}
    if not df.empty:
        for _, row in df.iterrows():
            db_strategies[row['strategy_name']] = row['count']
    
    # 从文件系统获取可用的 runner 文件
    backtest_dir = os.path.join(project_root, 'backtest')
    file_strategies = {}
    
    if os.path.exists(backtest_dir):
        for file in os.listdir(backtest_dir):
            if file.startswith('runner_') and file.endswith('.py') and not file.startswith('temp_'):
                # 提取策略名称（去掉 runner_ 前缀和 .py 后缀）
                strategy_key = file.replace('runner_', '').replace('.py', '')
                file_strategies[strategy_key] = file
    
    # 构建返回结果
    strategies = []
    
    # 1. 先添加数据库中有记录的策略
    for strategy_name, count in db_strategies.items():
        # 标准化策略名称（兼容旧名称）
        runner_key = normalize_strategy_name(strategy_name)
        runner_file = file_strategies.get(runner_key)
        
        strategies.append({
            'id': strategy_name,
            'name': f'{strategy_name} ({count}次)',
            'runner_key': runner_key,
            'file': runner_file or f'runner_{runner_key}.py',
            'description': f'{strategy_name} - 已运行{count}次',
            'has_runner': runner_file is not None,
            'count': count
        })
    
    # 2. 添加文件系统中有但数据库中没有的策略
    for strategy_key, file_name in file_strategies.items():
        # 检查是否已经添加过
        if strategy_key not in db_strategies:
            # 将 runner key 转换为友好的显示名称
            display_name = strategy_key.replace('_', ' ').title()
            
            strategies.append({
                'id': strategy_key,
                'name': display_name,
                'runner_key': strategy_key,
                'file': file_name,
                'description': f'{display_name} 策略',
                'has_runner': True,
                'count': 0
            })
    
    return jsonify(strategies)


# ==================== 路由：获取策略默认参数 ====================
@app.route('/api/strategy/<strategy_name>/params')
def get_strategy_params(strategy_name):
    """获取策略的默认参数配置（从数据库最近一次运行或默认值）"""
    # 先尝试从数据库获取最近一次的参数
    conn = sqlite3.connect(db_path)
    query = '''
        SELECT strategy_params, symbol, start_date, end_date, init_balance
        FROM backtest_results
        WHERE strategy_name = ?
        ORDER BY created_at DESC
        LIMIT 1
    '''
    df = pd.read_sql_query(query, conn, params=(strategy_name,))
    conn.close()
    
    if not df.empty:
        # 从数据库获取参数
        try:
            params = json.loads(df.iloc[0]['strategy_params'])
            # 提取第一个品种的参数
            if params and isinstance(params, dict):
                first_symbol = list(params.keys())[0]
                symbol_params = params[first_symbol]
                
                # 添加基本参数
                result = {
                    'symbol': first_symbol,
                    'exchange': symbol_params.get('exchange', 'SHFE'),
                    'start_date': df.iloc[0]['start_date'],
                    'end_date': df.iloc[0]['end_date'],
                    'init_balance': int(df.iloc[0]['init_balance'])
                }
                
                # 添加策略特定参数
                for key, value in symbol_params.items():
                    if key not in ['exchange', 'i_symbol', 'm_symbol']:
                        result[key] = value
                
                return jsonify(result)
        except:
            pass
    
    # 如果数据库没有，返回默认参数
    default_params = {
        'symbol': 'rb',
        'exchange': 'SHFE',
        'start_date': '2024-01-01',
        'end_date': '2024-12-31',
        'init_balance': 10000000,
        'market_period': 1440
    }
    
    return jsonify(default_params)


# ==================== 路由：运行回测 ====================
@app.route('/api/backtest/run', methods=['POST'])
def run_backtest():
    """运行回测"""
    global current_process, current_output, current_webui_port
    
    data = request.json
    strategy_name = data.get('strategy_id') or data.get('strategy_name')  # 兼容两种参数名
    params = data.get('params', {})
    
    # 生成随机端口（10000-65535）
    import random
    webui_port = random.randint(10000, 65535)
    current_webui_port = webui_port
    
    # 生成临时运行脚本
    script_path = generate_runner_script(strategy_name, params, webui_port)
    
    if not script_path:
        return jsonify({'success': False, 'message': '不支持的策略'}), 400
    
    # 清空输出缓存
    current_output = []
    
    # 在后台线程中运行回测
    def run_in_background():
        global current_process, current_output
        try:
            current_process = subprocess.Popen(
                ['python', script_path],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1
            )
            
            # 实时读取输出
            for line in current_process.stdout:
                current_output.append(line)
                if len(current_output) > 1000:  # 限制缓存大小
                    current_output.pop(0)
            
            current_process.wait()
        except Exception as e:
            current_output.append(f"错误: {str(e)}\n")
    
    thread = threading.Thread(target=run_in_background)
    thread.daemon = True
    thread.start()
    
    return jsonify({
        'success': True, 
        'message': '回测已启动',
        'webui_port': webui_port
    })


# ==================== 路由：获取回测输出 ====================
@app.route('/api/backtest/output')
def get_backtest_output():
    """获取回测实时输出"""
    global current_output, current_webui_port
    return jsonify({
        'output': ''.join(current_output[-100:]),
        'webui_port': current_webui_port
    })  # 返回最后100行


# ==================== 路由：停止回测 ====================
@app.route('/api/backtest/stop', methods=['POST'])
def stop_backtest():
    """停止当前运行的回测"""
    global current_process
    
    if current_process and current_process.poll() is None:
        current_process.terminate()
        return jsonify({'success': True, 'message': '回测已停止'})
    
    return jsonify({'success': False, 'message': '没有正在运行的回测'})


# ==================== 路由：获取历史回测结果 ====================
@app.route('/api/backtest/history')
def get_backtest_history():
    """获取历史回测结果列表"""
    strategy_name = request.args.get('strategy_name', '')
    limit = int(request.args.get('limit', 50))
    
    conn = sqlite3.connect(db_path)
    
    query = '''
        SELECT id, strategy_name, symbol, start_date, end_date, 
               init_balance, final_balance, total_return, annual_yield,
               sharpe_ratio, max_drawdown, winning_rate, open_times, close_times,
               test_timestamp, created_at
        FROM backtest_results
    '''
    
    if strategy_name:
        query += f" WHERE strategy_name LIKE '%{strategy_name}%'"
    
    query += ' ORDER BY created_at DESC LIMIT ?'
    
    df = pd.read_sql_query(query, conn, params=(limit,))
    conn.close()
    
    # 转换为字典列表
    results = df.to_dict('records')
    
    return jsonify(results)


# ==================== 路由：获取回测详情 ====================
@app.route('/api/backtest/<int:backtest_id>')
def get_backtest_detail(backtest_id):
    """获取单个回测的详细信息"""
    conn = sqlite3.connect(db_path)
    
    # 获取基本信息
    result_query = 'SELECT * FROM backtest_results WHERE id = ?'
    result_df = pd.read_sql_query(result_query, conn, params=(backtest_id,))
    
    if result_df.empty:
        conn.close()
        return jsonify({'error': '回测结果不存在'}), 404
    
    result = result_df.to_dict('records')[0]
    
    # 获取账户记录（用于绘制权益曲线）
    account_query = '''
        SELECT trade_date, balance, available, float_profit, 
               close_profit, margin, risk_ratio
        FROM account_records
        WHERE backtest_id = ?
        ORDER BY trade_date
    '''
    account_df = pd.read_sql_query(account_query, conn, params=(backtest_id,))
    
    # 获取交易明细
    trade_query = '''
        SELECT trade_date, datetime, symbol, direction, offset, 
               volume, price, commission
        FROM trade_details
        WHERE backtest_id = ?
        ORDER BY trade_date, datetime
    '''
    trade_df = pd.read_sql_query(trade_query, conn, params=(backtest_id,))
    
    conn.close()
    
    return jsonify({
        'result': result,
        'account_records': account_df.to_dict('records'),
        'trades': trade_df.to_dict('records')
    })


# ==================== 路由：对比分析 ====================
@app.route('/api/backtest/compare')
def compare_backtests():
    """对比多个回测结果"""
    strategy_name = request.args.get('strategy_name', '')
    
    conn = sqlite3.connect(db_path)
    
    query = '''
        SELECT id, strategy_name, symbol, start_date, end_date,
               total_return, annual_yield, sharpe_ratio, max_drawdown,
               winning_rate, strategy_params, test_timestamp
        FROM backtest_results
    '''
    
    if strategy_name:
        query += f" WHERE strategy_name = '{strategy_name}'"
    
    query += ' ORDER BY test_timestamp DESC LIMIT 20'
    
    df = pd.read_sql_query(query, conn)
    conn.close()
    
    # 解析参数
    results = []
    for _, row in df.iterrows():
        try:
            params = json.loads(row['strategy_params']) if row['strategy_params'] else {}
        except:
            params = {}
        
        results.append({
            'id': row['id'],
            'strategy_name': row['strategy_name'],
            'symbol': row['symbol'],
            'total_return': row['total_return'],
            'annual_yield': row['annual_yield'],
            'sharpe_ratio': row['sharpe_ratio'],
            'max_drawdown': row['max_drawdown'],
            'winning_rate': row['winning_rate'],
            'params': params,
            'test_timestamp': row['test_timestamp']
        })
    
    return jsonify(results)


# ==================== 辅助函数：生成运行脚本 ====================
def normalize_strategy_name(strategy_name):
    """将旧的策略类名转换为新的 strategy_name"""
    # 旧名称到新名称的映射（兼容数据库中的旧记录）
    name_mapping = {
        'RSRStrategy': 'rsrs',
        'MultiSignalSVMStrategy': 'multi_signal_svm',
        'MultiSignalEnsembleStrategy': 'multi_signal_ensemble',
        'OcDiffStrategy': 'oc_diff',
        'MultiSignalStrategy': 'multi_signal',
        'EnhancedMultiSignalStrategy': 'enhanced_multi_signal',
        'Wq101MLStrategy': 'wq101_ml',
        'WQ101MLStrategy': 'wq101_ml',
        'Wq101ZscoreStrategy': 'wq101_zscore',
        'WQ101ZscoreStrategy': 'wq101_zscore'
    }
    
    # 如果是旧名称，转换为新名称
    if strategy_name in name_mapping:
        return name_mapping[strategy_name]
    
    # 如果已经是新名称，直接返回
    return strategy_name


def generate_runner_script(strategy_name, params, webui_port):
    """根据参数生成临时运行脚本"""
    # 标准化策略名称（兼容旧名称）
    strategy_name = normalize_strategy_name(strategy_name)
    
    template_file = f'runner_{strategy_name}.py'
    
    # 读取模板文件
    template_path = os.path.join('trading_system', 'backtest', template_file)
    
    if not os.path.exists(template_path):
        return None
    
    with open(template_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 替换 Web UI 端口
    # 查找并替换 web_gui 参数
    import re
    
    # 匹配 web_gui="http://127.0.0.1:端口号" 或 web_gui='http://127.0.0.1:端口号'
    pattern = r'web_gui=["\']http://127\.0\.0\.1:\d+["\']'
    replacement = f'web_gui="http://127.0.0.1:{webui_port}"'
    content = re.sub(pattern, replacement, content)
    
    # 替换日期参数（去掉前导零避免八进制错误）
    if 'start_date' in params:
        # 解析日期并去掉前导零
        date_parts = params["start_date"].split('-')
        year, month, day = int(date_parts[0]), int(date_parts[1]), int(date_parts[2])
        content = re.sub(
            r'start_date = date\(\d+, \d+, \d+\)',
            f'start_date = date({year}, {month}, {day})',
            content
        )
    
    if 'end_date' in params:
        # 解析日期并去掉前导零
        date_parts = params["end_date"].split('-')
        year, month, day = int(date_parts[0]), int(date_parts[1]), int(date_parts[2])
        content = re.sub(
            r'end_date = date\(\d+, \d+, \d+\)',
            f'end_date = date({year}, {month}, {day})',
            content
        )
    
    # 替换初始资金
    if 'init_balance' in params:
        content = re.sub(
            r'init_balance = \d+',
            f'init_balance = {params["init_balance"]}',
            content
        )
    
    # 替换品种和交易所
    if 'symbol' in params:
        content = re.sub(
            r'"[a-zA-Z]+": \{[\s\S]*?"exchange": "[A-Z]+"',
            f'"{params["symbol"]}": {{\n        "exchange": "{params.get("exchange", "SHFE")}"',
            content,
            count=1
        )
    
    # 生成临时文件
    temp_file = os.path.join('trading_system', 'backtest', f'temp_{strategy_name}_{datetime.now().strftime("%Y%m%d_%H%M%S")}.py')
    
    with open(temp_file, 'w', encoding='utf-8') as f:
        f.write(content)
    
    return temp_file


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=2025, debug=True)
